package com.leetcode.tree;

import com.leetcode.basic.TreeNode;

import java.util.*;

/**
 * @author Dennis Li
 * @date 2020/10/13 20:46
 */
public class HuffmanTree {

    public void res(int[] nums) {
        TreeNode root = buildTree(nums);
        int sum = countWeight(root);
        System.out.println(sum);
    }

    public TreeNode buildTree(int[] nums) {
        Queue<TreeNode> queue = new PriorityQueue<>(Comparator.comparingInt(o -> o.val));

        for (int i = 0; i < nums.length; i++) {
            int w = nums[i];
            TreeNode node = new TreeNode(w);
            queue.offer(node);
        }

        while (queue.size() > 1) {
            TreeNode left = queue.poll();
            TreeNode right = queue.poll();
            TreeNode root = new TreeNode(left.val + right.val);
            root.left = left;
            root.right = right;
            queue.offer(root);
        }

        return queue.poll();
    }

    private int countWeight(TreeNode root) {
        if (root == null) return 0;

        Queue<TreeNode> queue = new PriorityQueue<>(Comparator.comparingInt(o -> o.val));
        queue.offer(root);

        int depth = 0, res = 0;

        while (!queue.isEmpty()) {
            int size = queue.size();
            while (size-- > 0) {
                TreeNode node = queue.poll();
                res += depth * node.val;
                if (node.left != null) queue.offer(node.left);
                if (node.right != null) queue.offer(node.right);
            }
            depth++;
        }
        return res;
    }

}
